/* Copyright 2021 Kerim Yalcin
**
**   Licensed under the Apache License, Version 2.0 (the "License");
**   you may not use this file except in compliance with the License.
**   You may obtain a copy of the License at
**
**     http://www.apache.org/licenses/LICENSE-2.0
**
**   Unless required by applicable law or agreed to in writing, software
**   distributed under the License is distributed on an "AS IS" BASIS,
**   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
**   See the License for the specific language governing permissions and
**   limitations under the License.
*/

#include "AccessToken.hpp"
#include <iostream>

namespace Agx {

AccessToken::AccessToken() :
    Base {"AccessToken"},
    
    mToken {nullptr},
    mProcess {nullptr},
    mThread {nullptr},
    mInformationType {TOKEN_INFORMATION_CLASS::TokenUser},
    mAccess_desired {0},
    mQueryHasSuccess {false},
    mUserAccounts {nullptr}
{}

AccessToken::~AccessToken() {
    CloseHandle(mToken);
    delete mUserAccounts;
}

AccessToken& AccessToken::free() {
    CloseHandle(mToken);
    mToken = nullptr;
    delete mUserAccounts;
    mUserAccounts = nullptr;
    return *this;
}

AccessToken& AccessToken::init() {
    mProcess = nullptr;
    mThread = nullptr;
    mInformationType = TOKEN_INFORMATION_CLASS::TokenUser;
    mAccess_desired = 0;
    mQueryHasSuccess = false;
    return *this;
}

AccessToken& AccessToken::set_Process(const HANDLE& pProcess) {
    mProcess = pProcess;
    return *this;
}

AccessToken& AccessToken::set_Thread(const HANDLE& pThread) {
    mThread = pThread;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_AdjustDefault() {
    mAccess_desired |= TOKEN_ADJUST_DEFAULT;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_AdjustGroups() {
    mAccess_desired |= TOKEN_ADJUST_GROUPS;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_AdjustPrivileges() {
    mAccess_desired |= TOKEN_ADJUST_PRIVILEGES;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_AdjustSessionID() {
    mAccess_desired |= TOKEN_ADJUST_SESSIONID;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_AssignPrimary() {
    mAccess_desired |= TOKEN_ASSIGN_PRIMARY;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_Duplicate() {
    mAccess_desired |= TOKEN_DUPLICATE;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_Excute() {
    mAccess_desired |= TOKEN_EXECUTE;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_Impersonate() {
    mAccess_desired |= TOKEN_IMPERSONATE;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_Query() {
    mAccess_desired |= TOKEN_QUERY;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_QuerySource() {
    mAccess_desired |= TOKEN_QUERY_SOURCE;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_Read() {
    mAccess_desired |= TOKEN_READ;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_Write() {
    mAccess_desired |= TOKEN_WRITE;
    return *this;
}

AccessToken& AccessToken::enable_Access_desired_All() {
    mAccess_desired |= TOKEN_ALL_ACCESS;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_AdjustDefault() {
    mAccess_desired ^= TOKEN_ADJUST_DEFAULT;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_AdjustGroups() {
    mAccess_desired ^= TOKEN_ADJUST_GROUPS;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_AdjustPrivileges() {
    mAccess_desired ^= TOKEN_ADJUST_PRIVILEGES;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_AdjustSessionID() {
    mAccess_desired ^= TOKEN_ADJUST_SESSIONID;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_AssignPrimary() {
    mAccess_desired ^= TOKEN_ASSIGN_PRIMARY;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_Duplicate() {
    mAccess_desired ^= TOKEN_DUPLICATE;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_Excute() {
    mAccess_desired ^= TOKEN_EXECUTE;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_Impersonate() {
    mAccess_desired ^= TOKEN_IMPERSONATE;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_Query() {
    mAccess_desired ^= TOKEN_QUERY;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_QuerySource() {
    mAccess_desired ^= TOKEN_QUERY_SOURCE;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_Read() {
    mAccess_desired ^= TOKEN_READ;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_Write() {
    mAccess_desired ^= TOKEN_WRITE;
    return *this;
}

AccessToken& AccessToken::disable_Access_desired_All() {
    mAccess_desired ^= TOKEN_ALL_ACCESS;
    return *this;
}


AccessToken& AccessToken::set_TokenInformationType_toUserAccount() {
    mInformationType = TOKEN_INFORMATION_CLASS::TokenUser;
    return *this;
}

AccessToken& AccessToken::set_TokenInformationType_toGroupAccount() {
    mInformationType = TOKEN_INFORMATION_CLASS::TokenGroups;
    return *this;
}

AccessToken& AccessToken::set_TokenInformationType_toPrivileges() {
    mInformationType = TOKEN_INFORMATION_CLASS::TokenPrivileges;
    return *this;
}

AccessToken& AccessToken::set_TokenInformationType_toOwnerSecurityIdentifier() {
    mInformationType = TOKEN_INFORMATION_CLASS::TokenOwner;
    return *this;
}

const TOKEN_INFORMATION_CLASS& AccessToken::get_TokenInformationType() const {
    return mInformationType;
}

AccessToken& AccessToken::query_TokenInformation() {
    mQueryHasSuccess = false;
    if (!mToken) {
        return *this;
    }

    DWORD bufferSize = 0;
    DWORD bufferSize_required = 0;
    bool hasFailed = false;

    switch (mInformationType) {
        case TOKEN_INFORMATION_CLASS::TokenUser:
            if (mUserAccounts) {
                delete mUserAccounts;
                mUserAccounts = nullptr;
            }
        break;
        default:
        break;
    }

    GetTokenInformation(
                mToken,
                mInformationType,
                nullptr,
                bufferSize,
                &bufferSize_required
            );

    const int size = get_array_size(bufferSize_required, sizeof(TOKEN_USER));

    switch (mInformationType) {
        case TOKEN_INFORMATION_CLASS::TokenUser:

            mUserAccounts = new TOKEN_USER[size];
            bufferSize = size*sizeof(TOKEN_USER);
            hasFailed = (
                GetTokenInformation(
                    mToken,
                    mInformationType,
                    mUserAccounts,
                    bufferSize,
                    &bufferSize_required
                )
            ) ?
                true : false;
        break;
        default:
        break;
    }

    mQueryHasSuccess = true;
    return *this;
}

AccessToken& AccessToken::query_ProcessToken() {
    if (
        !OpenProcessToken(
            mProcess,
            mAccess_desired,
            &mToken
        )
    ) {
        mError.print_fromWin32("query_ProcessToken");
    }
    return *this;
}


const bool& AccessToken::has_Query_success() const {
    return mQueryHasSuccess;
}

}